
#include <openssl/objects.h>
#include "SecurityTokenKey.h"
#include "../crypt/RSAPrivateKey.h"
#include "../crypt/RSAPublicKey.h"
#include "../crypt/DSAPrivateKey.h"
#include "../crypt/DSAPublicKey.h"
#include "../crypt/ECDSAPrivateKey.h"
#include "../crypt/ECDSAPublicKey.h"
#include "../utils/Hash.h"

/** Constructor for the security token with the names of appropriate
 * keys from the database.
 * 
 * @param isgnname The name of the signing key to load from the
 * database.
 * 
 * @param ichkname The name of the checking key to load from the
 * database.
 * 
 * @param ipassword The password to use to decrpyt the keys from the
 * database.  In this case, NULL means to use the default per-table
 * WIP password.
 * 
 * @param ikdb The key database. */
SecurityTokenKey::SecurityTokenKey(const char *isgnname, const char *ichkname,
   const char *ipassword, KeyDB *ikdb)
{
   tried = 0;
   sgnkey = Key::getSigningKey(ikdb, ipassword, isgnname);
   chkkey = Key::getCheckingKey(ikdb, ipassword, ichkname);
   myself = this;
}

/** Constructor for the security token with a key.
 * 
 * @param isgnkey Directly specified signing key for the SSH
 * challenge.
 * 
 * @param ichkkey Directly specified checking key for the SSH
 * challenge. */
SecurityTokenKey::SecurityTokenKey(const SigningKey &isgnkey, const CheckingKey &ichkkey)
{
   sgnkey = NULL;
   chkkey = NULL;
   switch (isgnkey.getKeyType())
   {
      case AbstractKey::KEY_TYPE_PRIVATE_RSA:
      {
         sgnkey = (SigningKey *)new RSAPrivateKey((const RSAPrivateKey &)(isgnkey));
         break;
      }
      case AbstractKey::KEY_TYPE_PRIVATE_DSA:
      {
         sgnkey = (SigningKey *)new DSAPrivateKey((const DSAPrivateKey &)(isgnkey));
         break;
      }
      case AbstractKey::KEY_TYPE_PRIVATE_ECDSA:
      {
         sgnkey = (SigningKey *)new ECDSAPrivateKey((const ECDSAPrivateKey &)(isgnkey));
         break;
      }
   }
   switch (ichkkey.getKeyType())
   {
      case AbstractKey::KEY_TYPE_PUBLIC_RSA:
      {
         chkkey = (CheckingKey *)new RSAPublicKey((const RSAPublicKey &)(ichkkey));
         break;
      }
      case AbstractKey::KEY_TYPE_PUBLIC_DSA:
      {
         chkkey = (CheckingKey *)new DSAPublicKey((const DSAPublicKey &)(ichkkey));
         break;
      }
      case AbstractKey::KEY_TYPE_PUBLIC_ECDSA:
      {
         chkkey = (CheckingKey *)new ECDSAPublicKey((const ECDSAPublicKey &)(ichkkey));
         break;
      }
   }
   myself = this;
}

/** Destructor for the key-based security token. */
SecurityTokenKey::~SecurityTokenKey()
{
   if (sgnkey) delete sgnkey;
   sgnkey = NULL;
   if (chkkey) delete chkkey;
   chkkey = NULL;
}

/** Callback for the key signing.
 * 
 * @param session The SSH session to authenticate.
 * 
 * @param[out] sig The signature that is produce by the action of our
 * signing key on the data presented to us to sign by the challenge.
 * 
 * @param[out] sig_len The length of the signature.
 * 
 * @param data The data that we are required to sign.
 * 
 * @param data_len The lenth of the data to sign.
 * 
 * @param abstract A user pointer (in this case the pointer to the
 * instance of the class that this was called from).
 * 
 * @return Zero on success, negative on failure. */
int SecurityTokenKey::keyCallback(LIBSSH2_SESSION *session,
   unsigned char **sig, size_t *sig_len, const unsigned char *data,
   size_t data_len, void **abstract)
{
   SecurityTokenKey *kthis = (SecurityTokenKey *)(*abstract);
   int length_sha1 = Hash::getBinarySHA1Hash(NULL, 0, NULL);
   char *binary_sha1 = (char *)malloc(sizeof(char)*length_sha1);
   Hash::getBinarySHA1Hash((const char *)data, data_len, (unsigned char *)binary_sha1);
   int retval = 0;
   unsigned int signature_size = 0;
   unsigned char *signature_blob = NULL;
   switch (kthis->sgnkey->getKeyType())
   {
      case AbstractKey::KEY_TYPE_PRIVATE_RSA:
      {
         RSA *rsa_n = NULL;
         ((RSAPrivateKey *)kthis->sgnkey)->getSSLKeyComponents(&rsa_n);
         signature_size = RSA_size(rsa_n);
         signature_blob = (unsigned char *)malloc(signature_size);
         if (signature_blob == NULL)
         {
            errprint("cannot allocate RSA signature");
            free(binary_sha1);
            RSA_free(rsa_n);
            return -1;
         }
         retval = RSA_sign(NID_sha1, (const unsigned char *)binary_sha1,
            length_sha1, signature_blob, &signature_size, rsa_n);
         free(binary_sha1);
         RSA_free(rsa_n);
         if (retval != 1)
            errprint("RSA signing failed");
         break;
      }
      case AbstractKey::KEY_TYPE_PRIVATE_DSA:
      {
         DSA *dsa_n = NULL;
         ((DSAPrivateKey *)kthis->sgnkey)->getSSLKeyComponents(&dsa_n);
         int r_size = 20;
         int s_size = 20;
         signature_size = r_size + s_size;
         signature_blob = (unsigned char *)calloc(signature_size, sizeof(unsigned char));
         DSA_SIG *dsasig = DSA_do_sign((const unsigned char *)binary_sha1,
            length_sha1, dsa_n);
         if ((!dsasig) || (BN_num_bytes(dsasig->r) > r_size) || (BN_num_bytes(dsasig->s) > s_size))
         {
            if (BN_num_bytes(dsasig->r) > r_size) errprint("r component too big");
            if (BN_num_bytes(dsasig->s) > s_size) errprint("s component too big");
            free(binary_sha1);
            DSA_free(dsa_n);
            if (!dsasig)
               errprint("cannot make DSA signature");
            else
               DSA_SIG_free(dsasig);
            return -1;
         }
         // right aligned within the cells
         BN_bn2bin(dsasig->r, &(signature_blob[r_size - BN_num_bytes(dsasig->r)]));
         BN_bn2bin(dsasig->s, &(signature_blob[r_size + (s_size - BN_num_bytes(dsasig->s))]));
         DSA_SIG_free(dsasig);
         free(binary_sha1);
         DSA_free(dsa_n);
         retval = 1;
         break;
      }
      case AbstractKey::KEY_TYPE_PRIVATE_ECDSA:
      {
         free(binary_sha1);
         break;
      }
      default:
      {
         errprint("unknown signing key type");
         free(binary_sha1);
         return -1;
         break;
      }
   }
   if (!retval)
   {
      if (signature_blob) free(signature_blob);
      errprint("signing failed");
      return -1;
   }
   *sig = signature_blob;
   *sig_len = signature_size;
   return 0;
}

/** SSH session authentication call.
 * 
 * @param tsession The SSH session to authenticate.
 * 
 * @param username The username to use with the key for SSH
 * authentication.
 * 
 * @return Zero on success, negative on failure. */
int SecurityTokenKey::authenticateSSHSession(LIBSSH2_SESSION *tsession, const char *username)
{
   void **abstract_pointer = libssh2_session_abstract(tsession);
   *abstract_pointer = (void *)this;
   char *datakeyblob = NULL;
   int datakeybloblen = 0;
   switch (chkkey->getKeyType())
   {
      case AbstractKey::KEY_TYPE_PUBLIC_RSA:
      {
         ((RSAPublicKey *)chkkey)->getSSHKeyBlob(&datakeyblob, &datakeybloblen);
         break;
      }
      case AbstractKey::KEY_TYPE_PUBLIC_DSA:
      {
         ((DSAPublicKey *)chkkey)->getSSHKeyBlob(&datakeyblob, &datakeybloblen);
         break;
      }
      case AbstractKey::KEY_TYPE_PUBLIC_ECDSA:
      {
         ((ECDSAPublicKey *)chkkey)->getSSHKeyBlob(&datakeyblob, &datakeybloblen);
         break;
      }
   }
   int retval = libssh2_userauth_publickey(tsession, username,
      (const unsigned char *)datakeyblob, datakeybloblen, getKeyCallback(), (void **)(&myself));
   free(datakeyblob);
   return retval;
}
